from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.tools import tool

from application.prompt.prompt_template_service import PromptTemplateService
from domain.model.register import CHAT_MODEL_MANAGER
from infrastructure.constant import file, models
from infrastructure.utils import file as file_utils
from infrastructure.utils import text, env


class EntityGoCodeGenerator:

    @staticmethod
    def generate(query: str, model_name: str = models.DEFAULT_MODEL) -> str:
        # 启动观察
        env.set_langsmith()
        # 得到聊天模型
        chat_model = CHAT_MODEL_MANAGER.get_model_instance_by_name(model_name)
        # 构造prompt
        prompt_value = PromptTemplateService.from_message_list_with_system_message(
            messages=[{'role': 'user', 'content': query}]
        )
        # 同步调用模型
        result_message = chat_model.invoke(input=prompt_value)
        content = result_message.content
        # 去除深度思考tag
        if "<think>" in content:
            content = text.remove_think_blocks(content).strip()
        # 去除go-markdown-code包裹
        if "```go" in content:
            content = text.extract_md_go_code_blocks(content).strip()
        # 保持生成内容到文件
        file_utils.save_text_content(file.GENERATE_GO_CODE_PATH, file.FILE_NAME_OF_ENTITY, content)
        return content

    @staticmethod
    def agent_generate(query: str, model_name: str = models.DEFAULT_MODEL) -> str:
        # 启动观察
        env.set_langsmith()
        # 得到聊天模型
        chat_model = CHAT_MODEL_MANAGER.get_model_instance_by_name(model_name)
        # 构造prompt
        prompt = ChatPromptTemplate.from_messages(
            messages=[
                ("system", PromptTemplateService.get_string_prompt_with_system_message()),
                ("human", "{{input}}"),
                ("placeholder", "{agent_scratchpad}"),
            ],
            template_format='jinja2'
        )
        # 加载可用工具
        tools = [ConfigTypeInterfaceTool, ITemplateEntityInterfaceTool, FrontCodeInstanceTool]
        agent = create_tool_calling_agent(chat_model, tools, prompt)
        agent_executor = AgentExecutor(agent=agent, tools=tools)
        # 调用agent
        result_message = agent_executor.invoke({'input': query})
        content = result_message.get('output')
        # 去除深度思考tag
        if "<think>" in content:
            content = text.remove_think_blocks(content).strip()
        # 去除go-markdown-code包裹
        if "```go" in content:
            content = text.extract_md_go_code_blocks(content).strip()
        # 保持生成内容到文件
        file_utils.save_text_content(file.GENERATE_GO_CODE_PATH, file.FILE_NAME_OF_ENTITY, content)
        return content


@tool
def ConfigTypeInterfaceTool():
    """该工具用于获取 ConfigType 接口的具体定义内容"""
    return """// ConfigType v3模板配置 interface
type ConfigType interface {
	ToJson() (string, error)                  // 转json
	ToEntity(data string) (ConfigType, error) // 转实体
}
    """


@tool
def ITemplateEntityInterfaceTool():
    """该工具用于获取 ITemplateEntity 接口的具体定义内容"""
    return """// ITemplateEntity v3 模板实体接口定义
type ITemplateEntity interface {
	// GetTemplateType tqt定义的唯一key
	GetTemplateType() EnumTemplateType
	// GetBigType 服务端定义的活动大类
	GetBigType() int
	// GetActType  服务端定义的活动小类
	GetActType() int
	// GetTemplateName 模板名称
	GetTemplateName() string
	// GetTemplateKey 模板key
	GetTemplateKey() TemplateKey
	// GetZoneSize n
	GetZoneSize() int
	// GetVersion 模板版本
	GetVersion() Version
	// GetStartTime 开始时间
	GetStartTime() string
	// GetEndTime 结束时间
	GetEndTime() string
	// GetConfig 模板奖励配置
	GetConfig() ConfigType
}
    """


@tool
def FrontCodeInstanceTool():
    """该工具用于获取 FrontCodeInstance 接口的具体定义内容"""
    return """// FrontCodeInstance 前端配置
type FrontCodeInstance interface {
	FrontCode(ctx context.Context) (fc excel_reader.FrontCode, err error)
}
    """
